import sys
import time
from PyQt5 import QtGui
sys.path.append(r"C:\Python Projects\PycharmProjects\界面")
from UI.form_deal_data import Ui_Form_deal_data
from PyQt5.Qt import *
import cgitb
import json
import os
import warnings
import numpy as np
import struct
warnings.filterwarnings('ignore')
cgitb.enable(format='txt')


class DealDataForm(QWidget, Ui_Form_deal_data):

    def __init__(self, parent=None):
        super(DealDataForm, self).__init__(parent)
        self.setupUi(self)
        with open("qss/deal_data_win.css", "r", encoding='utf8') as f:
            qss = f.read()
        self.setStyleSheet(qss)
        self.setWindowTitle("数据处理")
        self.pushButton_getoridata_path.setStyleSheet("min-width: 25px; height:10px")
        self.pushButton_getsavedata_path.setStyleSheet("min-width: 25px;height:10px")
        # self.setFixedSize(920, 750)
        self.lineEdit_oridata_path.setReadOnly(True)
        self.lineEdit_data_save_path.setReadOnly(True)
        self.spinBox_genDataCnt.setMaximum(1000000000)
        self.paused = False
        self.running = False
        self.lineEdit_oridata_path.cursorPositionChanged.connect(lambda :self.lineEdit_oridata_path.setCursorPosition(0))
        self.lineEdit_data_save_path.cursorPositionChanged.connect(lambda :self.lineEdit_data_save_path.setCursorPosition(0))
        self.test()


    def test(self):
        self.lineEdit_oridata_path.setText(os.path.abspath(r"C:\e\data\DATA5yi.bin"))
        self.lineEdit_data_save_path.setText(os.path.abspath(r"C:\Users\weixuan\Desktop\test\train_dta"))


    def closeEvent(self, a0: QtGui.QCloseEvent) -> None:
        try:
            if self.running:
                self.thread.terminate()
            a0.setAccepted(True)
        except Exception as e:
            print(e)

    @pyqtSlot()
    def on_pushButton_recover_clicked(self):
        if self.paused and self.running:
            self.thread.mutex.unlock()
            self.paused = False
            self.thread.timer.start(1000)
        print("恢复")

    @pyqtSlot()
    def on_pushButton_pause_clicked(self):
        if not self.paused and self.running:
            self.thread.mutex.lock()
            self.paused = True
            self.thread.timer.stop()
        print("暂停")

    def printToTextBrowser(self, text, color='white', font_size='3', font_family='Times New Roman'):
        content = f"""
        <font size={font_size} color={color} style= "font-family:{font_family}">{text}</font>
        """
        self.textBrowser_output.append(content)


    def showFinished(self, s):
        self.printToTextBrowser(s)


    def showProcess(self,cnt):
        print(cnt)
        self.printToTextBrowser("已处理%d/%d个" % (cnt, self.train_cnt) )

    @pyqtSlot()
    def on_pushButton_deal_data_clicked(self):
        if self.running:
            return
        try:
            nfft = int(sys.argv[1])
            nperseg = int(sys.argv[2])
            noverlap = int(sys.argv[3])
            sample_freq = float(sys.argv[4])
            time_segment = int(sys.argv[5])
        except Exception as e:
            print(e)
        if not self.running:
            train_cnt = self.spinBox_genDataCnt.value()
            if train_cnt == 0:
                QMessageBox.critical(self, "错误提示", "训练数据个数不能为0")
                return
            self.train_cnt = train_cnt
            bin_data_path = self.lineEdit_oridata_path.text()
            if bin_data_path == "":
                QMessageBox.critical(self, "错误提示", "请输入IQ数据路径")
                return
            save_path = self.lineEdit_data_save_path.text()
            if save_path == "":
                QMessageBox.critical(self, "错误提示", "请输入数据保存路径")
                return
            if not os.path.exists(save_path):
                os.makedirs(save_path)

            self.thread = BinFileReader(self, train_cnt = train_cnt,
                                        bin_file_path= bin_data_path,
                                        save_path=save_path,
                                        noverlap=noverlap,
                                        time_segment= time_segment,
                                        nfft = nfft,
                                        Fs = sample_freq,
                                        nperseg= nperseg)
            self.thread.requestShowPromption[int].connect(self.showProcess)
            self.thread.requestShowFinished[str].connect(self.showFinished)

            self.thread.start()
            self.running = True

        self.printToTextBrowser("数据处理使用相关参数如下:")
        s = self.GetHtml(['窗长(nperseg)', '窗函数重叠数(noverlap)', 'FFT长度(nfft)', '时频图时间维度长度', '采样频率(MHz)', '单张时频图维度'],
                    [[nperseg, noverlap, nfft, time_segment, sample_freq, "(%d, %d)" % (time_segment, nfft)]]
                   )
        self.textBrowser_output.append(s)

    @pyqtSlot()
    def on_pushButton_getoridata_path_clicked(self):
        path, _ = QFileDialog.getOpenFileName(self, "IQ数据路径选取", ".", "Bin Files(*.bin)")
        if path == "":
            return
        self.lineEdit_oridata_path.setText(os.path.abspath(path))
        self.lineEdit_oridata_path.setCursorPosition(0)

    @pyqtSlot()
    def on_pushButton_getsavedata_path_clicked(self):
        path = QFileDialog.getExistingDirectory(self, "数据保存路径选取", ".", QFileDialog.ShowDirsOnly)
        if path == "":
            return
        self.lineEdit_data_save_path.setText(os.path.abspath(path))
        self.lineEdit_data_save_path.setCursorPosition(0)

flag = False
index = 0
# 读取文件线程
class BinFileReader(QThread):
    requestShowPromption = pyqtSignal(int)
    requestShowFinished = pyqtSignal(str)

    def __init__(self, parent, train_cnt,bin_file_path,save_path ,noverlap, time_segment=10, nfft=2048, Fs=100e6, nperseg=None):
        super(BinFileReader, self).__init__(parent)
        self.mutex = QMutex()
        self.nperseg = nfft if nperseg is None else nperseg
        self.train_cnt = train_cnt
        self.savepath = save_path
        self.parent = parent
        self.bin_file_path = bin_file_path
        self.fp = None
        self.nfft = nfft
        self.Fs = Fs
        self.noverlap = noverlap
        self.width = nfft
        self.time_segment = time_segment
        self.iq_complex_cnt = self.getReadLength(self.time_segment, self.nfft, self.noverlap)
        self.timer = QTimer(self)
        self.timer.timeout.connect(self.showProcess)

    def showProcess(self):
        print("showProcess")
        self.requestShowPromption.emit(index)

    def openFile(self):
        try:
            self.fp = open(self.bin_file_path, "rb")  # 以二进制只读方式打开文件
            return True
        except OSError:
            return False

    def getReadLength(self, time_col, window_length, noverlap):
        return (time_col - 3) * (window_length - noverlap) + window_length


    def run(self):
        print("进入文件读取线程")
        try:
            ret = self.openFile()
            if not ret:
                QMessageBox.critical(self, "错误提示","打开文件失败")
                return
        except:
            QMessageBox.critical(self, "错误提示","打开文件失败")
            print("打开测试文件失败")
            return
        # 读取文件
        if self.isInterruptionRequested():
            self.parent.running = False
            return
        train_data = np.zeros((self.train_cnt, self.time_segment, self.nfft), dtype=np.float) # 预先分配内存

        iq_complex_cnt = self.iq_complex_cnt
        buffer_num = 1
        read_num = iq_complex_cnt * 2  # 生成一张时頻图需要的IQ点数(非复数)
        read_bin_cnt = read_num * 2 * buffer_num

        I_index = np.arange(0, read_num, 2)
        Q_index = np.arange(1, read_num, 2)
        nfft = self.nfft
        nperseg = self.nperseg
        window = scipy.signal.get_window('hamming', Nx=nperseg)
        Fs = self.Fs
        noverlap = self.noverlap
        print("数据加载完毕")

        def IQ_to_PSD(IQ):
            global flag, index
            length = IQ.shape[0]
            for i in range(length):
                if self.isInterruptionRequested():
                    return
                I = IQ[i][I_index]
                Q = IQ[i][Q_index]
                IQ_complex = np.zeros([iq_complex_cnt], dtype=np.complex)
                IQ_complex.real = I
                IQ_complex.imag = Q
                F, T, Zxx = stft(IQ_complex, fs=Fs, noverlap=noverlap, nfft=nfft, window=window,
                                 return_onesided=False,
                                 nperseg=nperseg)
                s: np.ndarray = np.abs(Zxx) * 2 / nfft
                # s = (s - s.min()) / (s.max() - s.min())
                s = s / s.std()

                train_data[index] = s.T
                index += 1
                if not flag:
                    print("生成时频数据的维度", s.T.shape)
                    flag = True

        self.timer.start(1000)
        while index < self.train_cnt:
            try:
                self.mutex.lock()
                if self.isInterruptionRequested():
                    self.parent.running = False
                    print("退出程序")
                    self.timer.stop()
                    return
                content = self.fp.read(read_bin_cnt)
                IQ = struct.unpack('h' * read_num * buffer_num, content)
                IQ = np.array(IQ).reshape([-1, read_num])
                IQ_to_PSD(IQ)  # IQ数据转PSD并保存到文件队列
                self.mutex.unlock()
            except Exception as e:
                print(e)
        a, b, c = train_data.shape
        if index >= self.train_cnt:
            self.requestShowFinished.emit("保存文件中...")
            self.timer.stop()
            np.save("%s/train_data(%dx%dx%d).npy" % (self.savepath, a, b, c), train_data)
            self.requestShowFinished.emit("保存文件完毕!")
            self.parent.running = False
            self.timer.stop()














if __name__ == '__main__':
    app = QApplication(sys.argv)
    win = DealDataForm()
    win.show()
    sys.exit(app.exec_())









